% Code by Kaveh Zamani 
% Civil and Environmental Engineering Department 
% University of California at Davis
% !> Calculates the diffusive portion of the constituent transport.
% !> It contains an explicit version of the diffusion operator and a general (involving all
% !> potential cases) diffusion operator as well, with a coefficient theta_stm for 
% !> selecting the level of implicitness. (theta_stm=0.5 is Crank Nicolson.).
% !> The matrix is solved via a tri-diagonal solver.  !
% !> The algoritm looks like this:
% !>   - This creates the diffusive fluxes sends them for modification for boundaries and then differences the fluxes to get the operator d/dx(Ad/dx). 
% !>         - Calculate interior and boundary fluxes 
% !>   - Construct right hand side with neumann boundary condition imposed
% !>   - Construct diffusion coefficeint matrix with neumann boundary condition imposed
% !>   - Solve the system

function [conc]=diffiuse(conc_prev,area,area_prev,area_lo,area_hi,...
                         area_lo_prev,area_hi_prev,disp_coef_lo, ...
                         disp_coef_hi,disp_coef_lo_prev,disp_coef_hi_prev,...
                         ncell,nvar,time_new,theta_stm,dt,dx);
                     
%% arguments                      
% (out)conc(ncell,nvar)              !< Concentration at new time
% (in) conc_prev(ncell,nvar)         !< Concentration at old time
% (in) area (ncell)                  !< Cell-centered area at new time
% (in) area_prev (ncell)             !< Cell-centered area at old time
% (in) area_lo (ncell)               !< Low side area centered in time
% (in) area_hi (ncell)               !< High side area centered in time 
% (in) area_lo_prev (ncell)          !< Low side area centered at old time
% (in) area_hi_prev (ncell)          !< High side area centered at old time 
% (in) disp_coef_lo (ncell,nvar)     !< Low side constituent dispersion coef. at new time
% (in) disp_coef_hi (ncell,nvar)     !< High side constituent dispersion coef. at new time
% (in) disp_coef_lo_prev(ncell,nvar) !< Low side constituent dispersion coef. at old time
% (in) disp_coef_hi_prev(ncell,nvar) !< High side constituent dispersion coef. at old time
% (in) time_new                      !< Instantaneous "new" time to which we are advancing
% (in) theta_stm                     !< Explicitness coefficient; 0 is explicit, 0.5 Crank-Nicolson, 1 full implicit  
%% local arguments
% explicit_diffuse_op(ncell,nvar)             !< Explicit diffusive operator
% down_diag(ncell,nvar)                       !< Values of the coefficients below diagonal in matrix
% center_diag(ncell,nvar)                     !< Values of the coefficients at the diagonal in matrix
% up_diag(ncell,nvar)                         !< Values of the coefficients above the diagonal in matrix
% right_hand_side(ncell,nvar)                 !< Right hand side vector
% time_prev                                   !< old time
%% comment
% This routine gives the effects of diffusion fluxes on each cell
% for a single time step (ie, explicit). This is needed for the advection step.
% It is also part of the right hand side of the implicit diffusion solver 
% matrix calculation. 

% Explicit diffusion operator construction. This creates the diffusive fluxes
% sends them for modification for boundaries and then differences the fluxes
% to get the operator d/dx(Ad/dx). 
% instantaneous function

%% body
time_prev = time_new - dt;

[explicit_diffuse_op]=explicit_diffusion_operator(conc_prev, ...
                                                  area_lo_prev,...
                                                  area_hi_prev, ...
                                                  disp_coef_lo_prev,... 
                                                  disp_coef_hi_prev,...
                                                  ncell, ...
                                                  nvar,...
                                                  time_prev,...
                                                  dx,...
                                                  dt);

[right_hand_side] = construct_right_hand_side(explicit_diffuse_op,   ... 
                                              area_prev,             ...
                                              area_lo_prev,          ...
                                              area_hi_prev,          ...
                                              disp_coef_lo_prev,     ...
                                              disp_coef_hi_prev,     ...
                                              conc_prev,             ...
                                              theta_stm,             ...
                                              ncell,                 ...
                                              time_prev,             ...
                                              nvar,                  ...  
                                              dx,                    ...
                                              dt);

% Construct the matrix for the diffusion solver
% without boundary condition modification or structure on interior of domain
 [center_diag,up_diag,down_diag]= ...
                        construct_diffusion_matrix(area,...
                                                   area_lo,          ...
                                                   area_hi,          ...
                                                   disp_coef_lo,     ...
                                                   disp_coef_hi,     ...
                                                   theta_stm,        ...
                                                   ncell,            ...
                                                   time_new,         ... 
                                                   nvar,             ... 
                                                   dx,               ...
                                                   dt);

% this must set outside
[center_diag,up_diag,down_diag,right_hand_side]= ...
                       boundary_diffusion_matrix(center_diag ,       ...
                                                 up_diag,            ...     
                                                 down_diag,          ...
                                                 right_hand_side,    ... 
                                                 conc_prev,          ...
                                                 explicit_diffuse_op,...
                                                 area,               ...
                                                 area_lo,            ...
                                                 area_hi,            ...          
                                                 disp_coef_lo,       ...
                                                 disp_coef_hi,       ...
                                                 theta_stm,          ...
                                                 ncell,              ...
                                                 time_new,           ... 
                                                 nvar,               ... 
                                                 dx,                 ...
                                                 dt);
                                             

%% assable matrix LHS from three arrays                                           
LHS = sparse(zeros(ncell));

for icell=1:ncell
   LHS(icell,icell)=center_diag(icell);  
end

for icell=1:ncell-1
   LHS(icell,icell+1) = up_diag(icell);
   LHS(icell+1,icell) = down_diag(icell+1);
end

%% Solve
conc= linsolvr(LHS,right_hand_side);                                                     

end      

%% inside functions and routines 
%  Construct the right hand side vector from previous step,
%  and impose Neumann boundary condition on it.
function [right_hand_side]= construct_right_hand_side( ...
                                          explicit_diffuse_op,   ... 
                                          area_prev,             ...
                                          area_lo_prev,          ...
                                          area_hi_prev,          ...
                                          disp_coef_lo_prev,     ...
                                          disp_coef_hi_prev,     ...
                                          conc_prev,             ...
                                          theta,                 ...
                                          ncell,                 ...
                                          time,                  ...
                                          nvar,                  ...  
                                          dx,                    ...
                                          dt);

%-----------------arguments  
                             
% ncell                                     !< Number of cells
% nvar                                      !< Number of variables
% (out)  right_hand_side(ncell,nvar)       !< The right hand side vector
% explicit_diffuse_op (ncell,nvar)  !< Explicit diffusion operator
% area_prev (ncell)                 !< Cell centered area at old time 
% conc_prev(ncell,nvar)             !< Concentration at old time
% area_lo_prev (ncell)              !< Low side area at old time
% area_hi_prev (ncell)              !< High side area at old time 
% disp_coef_lo_prev (ncell,nvar)    !< Low side constituent dispersion coef. at old time
% disp_coef_hi_prev (ncell,nvar)    !< High side constituent dispersion coef. at old time
% time                              !< Current time
% theta                             !< Explicitness coefficient; 0 is explicit, 0.5 Crank-Nicolson, 1 full implicit  
% dx                                !< Spatial step  
% dt                                !< Time step 
  
for ivar = 1:nvar
     right_hand_side(:,ivar) = area_prev(:)*conc_prev(:,ivar) ...
                              - (one-theta)*dt*explicit_diffuse_op(:,ivar); 
end 

end

% Construct the matrix for the diffusion solver
% without boundary condition modification or structure on interior of domain
function[center_diag,up_diag,down_diag]=construct_diffusion_matrix(   ...
                                                      area,             ...
                                                      area_lo,          ...
                                                      area_hi,          ...
                                                      disp_coef_lo,     ...
                                                      disp_coef_hi,     ...
                                                      theta_stm,        ...
                                                      ncell,            ...
                                                      time,             ... 
                                                      nvar,             ... 
                                                      dx,               ...
                                                      dt);
                          
% ---args                             
% ncell                            !< Number of cells
% nvar                             !< Number of variables
% (out) :: down_diag(ncell,nvar)   !< Values of the coefficients below diagonal in matrix
% (out) :: center_diag(ncell,nvar) !< Values of the coefficients at the diagonal in matrix
% (out) :: up_diag(ncell,nvar)     !< Values of the coefficients above the diagonal in matrix
% area (ncell)             !< Cell centered area at new time 
% area_lo(ncell)           !< Low side area at new time
% area_hi(ncell)           !< High side area at new time 
% disp_coef_lo(ncell)      !< Low side constituent dispersion coef. at new time
% disp_coef_hi(ncell)      !< High side constituent dispersion coef. at new time
% time                     !< Current time
% theta_stm                !< Explicitness coefficient; 0 is explicit, 0.5 Crank-Nicolson, 1 full implicit 
% dx                       !< Spatial step  
% dt                       !< Time step                                   
                              
up_diag(ncell,:) = LARGEREAL;
down_diag(1,:) = LARGEREAL;
dt_by_dxsq = dt/(dx*dx);  

for ivar = 1:nvar 
    for icell = 1:ncell
     down_diag(icell,ivar) = -theta_stm*dt_by_dxsq*area_lo(icell)*disp_coef_lo(icell); 
     center_diag(icell,ivar) = ...
                 area(icell)+ theta_stm*dt_by_dxsq*(area_hi(icell)*disp_coef_hi(icell) ...
                                      + area_lo(icell)*disp_coef_lo(icell));
     up_diag(icell,ivar)=-theta_stm*dt_by_dxsq*area_hi(icell)*disp_coef_hi(icell);             
    end 
end   

end 

%  Calculate the explicit diffusion operator for a moment in time.
%  Explicit diffusion operator construction. This creates the diffusive fluxes
%  sends them for modification for boundaries and then differences the fluxes
%  to get the operator d/dx(Ad/dx). 
function[explicit_diffuse_op]=explicit_diffusion_operator( ...
                                                          conc,            ...
                                                          area_lo,         ...
                                                          area_hi,         ...
                                                          disp_coef_lo,    ...  
                                                          disp_coef_hi,    ...
                                                          ncell,           ...
                                                          nvar,            ...
                                                          time,            ...
                                                          dx,              ...
                                                          dt)

% ---args
% ncell                                 !< Number of cells
% nvar                                  !< Number of variables
% (out)explicit_diffuse_op(ncell,nvar)  !< Explicit diffusion operator
% conc(ncell,nvar)               !< Concentration at old time
% area_lo(ncell)                 !< Low side area at old time
% area_hi(ncell)                 !< High side area at old time 
% disp_coef_lo(ncell)            !< Low side constituent dispersion coef. at old time
% disp_coef_hi(ncell)            !< High side constituent dispersion coef. at old time
% time                           !< Current time
% dx                             !< Spacial step  
% dt                             !< Time step

for ivar = 1:nvar
    diffusive_flux_lo(2:ncell,ivar) = ...
        -(area_lo(2:ncell)*disp_coef_lo(2:ncell)* ...
        (conc(2:ncell,ivar) - conc(1:(ncell-1),ivar)))/dx;
               
    diffusive_flux_hi(1:(ncell-1),ivar) = ...
        -(area_hi(1:(ncell-1))*disp_coef_hi(1:(ncell-1))* ...
           (conc(2:ncell,ivar) - conc(1:(ncell-1),ivar)))/dx;                    
end

diffusive_flux_hi(ncell,:) = 123456789;
diffusive_flux_lo(1,:) = 123456789;

[diffusive_flux_lo,diffusive_flux_hi]=boundary_diffusion_flux(conc,        ... 
                                                              area_lo,     ...
                                                              area_hi,     ...
                                                              disp_coef_lo,...  
                                                              disp_coef_hi,...
                                                              ncell,       ...
                                                              nvar,        ...
                                                              time,        ...
                                                              dx,          ...
                                                              dt)

          
    for ivar = 1:nvar
       explicit_diffuse_op(:,ivar) = ...
           (diffusive_flux_hi(:,ivar) - diffusive_flux_lo(:,ivar))/dx;
    end 

end  


   